const std = @import("std");
const log = std.log.scoped(.trap);
const csr = @import("riscv").csr;
const os = @import("os.zig");
const proc = os.proc;
const driver = os.driver;
const color = os.lib.color;

pub const clock = @import("trap/clock.zig");
pub const trampoline = @import("trap/trampoline.zig");
const TRAMPOLINE = os.mem.VirtAddr.TRAMPOLINE;

// in kernelvec.zig, calls kerneltrap().
const kernelvec = @import("trap/kernelvec.zig").kernelvec;

/// set up to take exceptions and traps while in the kernel.
pub fn inithart() void {
    csr.stvec.write(&kernelvec);
}

/// handle an interrupt, exception, or system call from user space.
fn usertrap() callconv(.c) noreturn {
    if (csr.sstatus.read().spp) {
        @panic("usertrap: not from user mode");
    }

    // send interrupts and exceptions to kerneltrap(), since we're now in the kernel.
    csr.stvec.write(&kernelvec);

    // save user program counter.
    const p = proc.current().?;
    const p_uproc = &p.ctx.uproc;
    const sepc = csr.sepc.read();
    p_uproc.trapframe.epc = @intFromPtr(sepc);

    sw: switch (handler()) {
        .irq => {},
        .timer, .syscall => |event| {
            if (p.isKilled()) p.exit(-1);

            if (event == .timer) {
                // give up the CPU if this is a timer interrupt.
                p.yield();
            } else {
                // sepc points to the ecall instruction, but we want to return to the next instruction.
                p_uproc.trapframe.epc += 4;

                // an interrupt will change sepc, scause, and sstatus, so enable only now that we're done
                // with those registers.
                os.arch.intr.enable();

                os.syscall.syscall();
            }
        },
        .page_fault => p_uproc.user.pageFault(csr.stval.read()) catch continue :sw .unknown,
        .unknown => {
            log.warn(
                "{s}:{} scause {f}{f}{f}, sepc {*}, stval=0x{x}",
                .{ p.name.get(), p.pid, color.yellow, csr.scause.read(), color.reset, sepc, csr.stval.read() },
            );

            p.setKilled();
        },
    }

    if (p.isKilled()) p.exit(-1);
    usertrapret();
}

/// set up trapframe and control registers for a return to user space
pub fn usertrapret() noreturn {
    const p = proc.current().?;
    const p_uproc = &p.ctx.uproc;

    // we're about to switch the destination of traps from kerneltrap() to usertrap(). because a trap
    // from kernel code to usertrap would be a disaster, turn off interrupts.
    _ = os.arch.intr.disable();

    // send syscalls, interrupts, and exceptions to uservec in trampoline
    const uservec_addr = @intFromPtr(&trampoline.uservec);
    const trampoline_addr = @intFromPtr(&trampoline._trampoline);

    const trampoline_uservec = TRAMPOLINE + (uservec_addr - trampoline_addr);
    csr.stvec.write(@ptrFromInt(trampoline_uservec));

    // set up trapframe values that uservec will need when the process next traps into the kernel.
    p_uproc.trapframe.satp = @bitCast(csr.satp.read()); // kernel page table
    p_uproc.trapframe.sp = p.kstack; // process's kernel stack
    p_uproc.trapframe.trap = &usertrap;
    p_uproc.trapframe.hartid = proc.Cpu.id(); // hartid for cpuid()

    // set up the registers that trampoline's sret will use to get to user space.

    // set S Previous Privilege mode to User.
    csr.sstatus.clear(.{ .spp = true }); // clear SPP to 0 for user mode
    csr.sstatus.set(.{ .spie = true }); // enable interrupts in user mode

    // set S Exception Program Counter to the saved user pc.
    csr.sepc.write(@ptrFromInt(p_uproc.trapframe.epc));

    // tell trampoline the user page table to switch to.
    const Satp = csr.satp.Type();
    const satp: Satp = .{
        .mode = if (os.is_64_bit) .sv39 else .sv32,
        .ppn = p_uproc.user.pt.addr.ppn,
    };

    // jump to userret in trampoline at the top of memory, which switches to the user page table,
    // restores user registers, and switches to user mode with sret.
    const userret_addr = @intFromPtr(&trampoline.userret);
    const trampoline_userret = TRAMPOLINE + (userret_addr - trampoline_addr);
    const userret_fn: *const fn (Satp) callconv(.c) noreturn = @ptrFromInt(trampoline_userret);
    userret_fn(satp);
}

/// interrupts and exceptions from kernel code go here via kernelvec, on whatever the current kernel
/// stack is.
export fn kerneltrap() void {
    const sepc = csr.sepc.read();
    const sstatus = csr.sstatus.read();

    if (!sstatus.spp) {
        @panic("kerneltrap: not from supervisor mode");
    }

    if (sstatus.sie) {
        @panic("kerneltrap: interrupts enabled");
    }

    switch (handler()) {
        .irq => {},
        .timer => if (proc.current()) |p| p.yield(), // give up the CPU if this is a timer interrupt.
        .syscall, .page_fault, .unknown => {
            // interrupt or trap from an unknown source
            log.warn(
                "scause {f}{f}{f}, sepc {*}, stval=0x{x}",
                .{ color.yellow, csr.scause.read(), color.reset, sepc, csr.stval.read() },
            );
            @panic("kerneltrap");
        },
    }

    // the yield() may have caused some traps to occur, so restore trap registers for use by kernelvec's
    // sepc instruction.
    csr.sepc.write(sepc);
    csr.sstatus.write(sstatus);
}

/// check if it's an external interrupt, timer interrupt or system call, and handle it.
fn handler() enum { irq, timer, syscall, page_fault, unknown } {
    const scause = csr.scause.read();
    if (scause.interrupt) {
        switch (scause.code.interrupt) {
            // this is a supervisor external interrupt, via PLIC.
            .@"Supervisor external interrupt" => {
                const target: driver.plic.Target = .{
                    .mode = .supervisor,
                    .hart = @intCast(proc.Cpu.id()),
                };

                // irq indicates which device interrupted.
                const irq = target.claim() orelse return .irq;

                switch (irq) {
                    .virtio0 => driver.virtio.disk.intr(),
                    .uart0 => driver.uart.intr(),
                    else => {
                        if (irq.pcieDev()) |dev| {
                            driver.pci.intr(dev);
                        } else {
                            log.warn("unexpected interrupt irq={}", .{@intFromEnum(irq)});
                        }
                    },
                }

                // the PLIC allows each device to raise at most one interrupt at a time; tell the PLIC
                // the device is now allowed to interrupt again.
                target.complete(irq);

                return .irq;
            },
            .@"Supervisor timer interrupt" => {
                clock.intr();
                return .timer;
            },
            else => {},
        }
    } else {
        switch (scause.code.exception) {
            .@"Environment call from U-mode" => return .syscall,
            .@"Load page fault", .@"Store/AMO page fault" => return .page_fault,
            else => {},
        }
    }

    return .unknown;
}
